/**
 * File: MultiFileFastqWriter.java
 * Created by: mhaimel
 * Created on: 20 Jul 2009
 * CVS:  $Id: MultiFileFastqWriter.java,v 1.3 2009/10/09 15:14:58 mhaimel Exp $
 */
package uk.ac.ebi.curtain.util.impl;

import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.OutputStream;
import java.io.PrintWriter;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.atomic.AtomicLong;
import java.util.concurrent.atomic.AtomicReference;

import org.apache.commons.lang.NotImplementedException;
import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;

import uk.ac.ebi.curtain.util.FileIndexReader;
import uk.ac.ebi.curtain.util.FileIndexWriter;
import uk.ac.ebi.curtain.util.file.FileAccessService;
import uk.ac.ebi.curtain.util.file.FileChannelAccess;
import uk.ac.ebi.curtain.utils.CurtainUncheckedException;
import uk.ac.ebi.curtain.utils.concurrent.UtilsCollect;
import uk.ac.ebi.curtain.utils.data.FileType;
import uk.ac.ebi.curtain.utils.data.IdentifierProcessor;
import uk.ac.ebi.curtain.utils.data.ReadType;
import uk.ac.ebi.curtain.utils.file.EntryMapper;
import uk.ac.ebi.curtain.utils.file.EntryMapperBean;
import uk.ac.ebi.curtain.utils.file.FileInfo;
import uk.ac.ebi.curtain.utils.file.WriteUtil;
import uk.ac.ebi.curtain.utils.io.impl.BufferedRafReader;
import uk.ac.ebi.curtain.utils.io.impl.FileIO;
import uk.ac.ebi.velvet.model.Strand;

/**
 * @author mhaimel
 *
 */
public class MultiFileFastqWriter implements FileIndexWriter {

	private Log log = LogFactory.getLog(this.getClass());
	private Map<FileInfo, Integer> inputFile2IdIdx = new HashMap<FileInfo, Integer>();
	private Map<Integer, FileInfo> inputId2FileIdx = new HashMap<Integer, FileInfo>();
	private Map<FileInfo, Integer> outputFile2IdIdx = new HashMap<FileInfo, Integer>();
	private Map<Integer, FileInfo> outputId2FileIdx = new HashMap<Integer, FileInfo>();
	private Map<String, Integer> category2OutfileIdx = new HashMap<String, Integer>();
	private Map<Integer, Integer> in2OutIdIdx = new HashMap<Integer, Integer>();
	private final FileAccessService fas;
	
	public MultiFileFastqWriter(FileAccessService fas) {
		this.fas = fas;
	}
	
	protected Log getLog() {
		return log;
	}
	
	private String _buildOutputIdx(FileInfo fi, Integer category){
		ReadType readType = fi.getReadType();
		if(readType.isPaired()){
			return readType.name()+"_"+category.toString();
		} else {
			return readType.name(); 
		}
	}
	
	public void addOutputFile(Integer category, FileInfo info){
		String idxStr = _buildOutputIdx(info, category);
		Integer idx = _getOrAddOutputFileIdx(info);
		category2OutfileIdx.put(idxStr, idx);
	}
	
	private void _addOutputFileId(FileInfo info, Integer id){
		outputFile2IdIdx.put(info, id);
		outputId2FileIdx.put(id, info);
	}
	
	private Integer _getOrAddOutputFileIdx(FileInfo out){
		Integer fId = _getOutputFileIdx(out);
		if(null == fId){
			fId = _nextOutputFildId();
			_addOutputFileId(out, fId);
		}
		return fId;
	}
	
	private Integer _getOutputFileIdx(FileInfo out){
		return outputFile2IdIdx.get(out);
	}
	
	/**
	 * {@inheritDoc}
	 */
	public void addContig(Integer inId, String key, Long pos) {
		_writeEntry(inId, pos);
	}
	
	private Integer _getOutput4Input(Integer inId){
		return in2OutIdIdx.get(inId);
	}
	
	private static class StatusMemento{
		private FileInfo inInfo;
		private Integer inId;
		private FileInfo outInfo;
		private Integer outId;
		private FileChannelAccess<BufferedRafReader> in;
		private FileChannelAccess<PrintWriter> out;
		
//		private void reset(){
//			this.inInfo = null;
//			this.inId = null;
//			this.outInfo = null;
//			this.outId = null;
//			this.in = null;
//			this.out = null;
//		}
	}
	
	private StatusMemento currMomento = new StatusMemento();
	
	private void _writeEntry(FileInfo inFile, Long inputPos){
		currMomento = _updateMomento(inFile, currMomento,true);
		_writeEntry(currMomento,inputPos);
	}
	
	private StatusMemento _updateMomento(Integer inId, StatusMemento prevMomento, boolean useOutChannel){
		StatusMemento momento = prevMomento;
		if(null == momento 
				|| momento.inInfo == null 
				|| momento.inId == null){
			momento = new StatusMemento();
			_updateInput(inId,momento);
			_updateOutput(momento,useOutChannel);
		} else {
			Integer prevOutId = momento.outId;
			
			/* update input file if needed */
			if(!momento.inId.equals(inId)){
				_updateInput(inId,momento);				
			}
			/* update output file if changed */
			if(null == prevOutId 
					|| !prevOutId.equals(momento.outId)){
				_updateOutput(momento,useOutChannel);				
			}
		}
		return momento;
	}
	
	private StatusMemento _updateMomento(FileInfo inFile, StatusMemento prevMomento, boolean useOutChannel){
		StatusMemento momento = prevMomento;
		if(null == momento 
				|| momento.inInfo == null 
				|| momento.inId == null){
			momento = new StatusMemento();
			_updateInput(inFile,momento);
			_updateOutput(momento,useOutChannel);
		} else {
			Integer prevOutId = momento.outId;
			
			/* update input file if needed */
			if(!momento.inInfo.equals(inFile)){
				_updateInput(inFile,momento);				
			}
			/* update output file if changed */
			if(null == prevOutId 
					|| !prevOutId.equals(momento.outId)){
				_updateOutput(momento,useOutChannel);				
			}
		}
		return momento;
	}
	
	private void _updateOutput(StatusMemento momento, boolean useChannel) {
		/* ensure it's closed */
		flushAndClose(momento.outInfo, momento.out);
		
		Integer outId = momento.outId;
		FileInfo outInfo = _getOutputFileInfo(outId);
		if(useChannel){
			FileChannelAccess<PrintWriter> out = getFileAccess().getPrintWriter(outInfo.getFile());
			momento.out = out;
		}
		momento.outInfo = outInfo;
	}

	private void _updateInput(FileInfo inFile, StatusMemento momento) {
		Integer inId = _getInputfileId(inFile);	
		_updateInput(inFile, inId, momento);
	}
	
	private void _updateInput(Integer inId, StatusMemento momento) {
		FileInfo inFile = _getInputFileInfo(inId);
		_updateInput(inFile, inId, momento);
	}
	
	private void _updateInput(FileInfo inFile, Integer inId, StatusMemento momento) {
		/* ensure it's closed */
		FileIO.closeQuietly(momento.in);
		Integer outId = _getOutput4Input(inId);
		FileChannelAccess<BufferedRafReader> in = getFileAccess().getBufferedRafReader(inFile.getFile());
		momento.inInfo = inFile;
		momento.inId = inId;
		momento.in = in;
		momento.outId = outId;
	}

	private Integer _getInputfileId(FileInfo inFile) {
		return inputFile2IdIdx.get(inFile);
	}

	private void _writeEntry(Integer inId, Long inputPos){
		currMomento = _updateMomento(inId,currMomento, true);
		_writeEntry(currMomento, inputPos);
	}
	
	private MyEntryMapper singleMapper = new MyEntryMapper();
	private void _writeEntry(StatusMemento momento, Long inputPos){
		try {
			FileType inType = momento.inInfo.getFileType();
			BufferedRafReader channel = momento.in.getChannel();
			channel.seek(inputPos);
			singleMapper.setChannelAccess(momento.out);
			singleMapper.setWriter(momento.outInfo.getFileType().getWriter());
			inType.getParser().parseNEntry(
					1,
					channel, 
					momento.inInfo, 
					singleMapper);			
		} catch (IOException e) {
			if(null != momento.in){ 
				momento.in.reportProblem();
			}
			_throwException("Problems writing entry: ",momento.inInfo,momento.outInfo,inputPos,e);
		}finally{
			FileIO.closeQuietly(momento.in);
			flushAndClose(momento.outInfo, momento.out);
		}
	}
	
	private void _writeBatch(Integer[] inArr, Long[] posArr){
		
		StatusMemento currMomento = new StatusMemento();
		long bufferSize = 1024*1024*4; // 4MB

		BufferedEntryMapper mapper = new BufferedEntryMapper();
		
		try{
			for(int i = 0; i < inArr.length; ++i){
				if(mapper.getByteCount() > bufferSize){
					_flush(mapper);
				}
				Integer currInId = inArr[i];
				Long currPos = posArr[i];
				Integer prevOutId = currMomento.outId;
				try {
					/* check if input file has changed */
					if(null == currMomento.inId 
							|| !currMomento.inId.equals(currInId)){
						/* clear previous resources */
//						FileIO.closeQuietly(in);
						_updateInput(currInId, currMomento);
//						/* get new files */
//						inFile = _getInputFileInfo(currInId);
//						outId = _getOutput4Input(currInId);
//						in = _getInputChannel(currMomento.inInfo);
					}
					
					/* check if output file has changed (mainly stays the same) */
					if(null == prevOutId || !currMomento.outId.equals(prevOutId)){
						/* clear previous resources */
						_updateOutput(currMomento, false);
						mapper.setWriter(currMomento.outInfo.getFileType().getWriter());
						mapper.setOutputId(currMomento.outId);
					}
					
					/* set input into correct position */
					BufferedRafReader channel = currMomento.in.getChannel();
					channel.seek(currPos);
					
					/* map entry */
					currMomento.inInfo.getFileType().getParser().parseNEntry(
							1, 
							channel, 
							currMomento.inInfo, 
							mapper);
				} catch (IOException e) {
					if(null != currMomento.in){ 
						currMomento.in.reportProblem();
					}
					_throwException("Problems writing balk entry: ",currMomento.inInfo,currMomento.outInfo,currPos,e);
				} /* DON'T close input or output - they might be reused in the next round */
			}

			if(mapper.getByteCount() > 0){
				_flush(mapper);
			}
		} finally{
			FileIO.closeQuietly(currMomento.in);
		}
	}

	private void _flush(BufferedEntryMapper mapper) {
		List<Integer> ids = new ArrayList<Integer>(mapper.getStreamMap().keySet());
		for(Integer outId : ids){
			ByteArrayOutputStream stream = mapper.getStreamMap().get(outId);
			PrintWriter pm = mapper.getPrintMap().get(outId);
			FileIO.closeQuietly(pm);
			if(stream.size() > 0){
				FileInfo outInfo = _getOutputFileInfo(outId);
				FileChannelAccess<OutputStream> out = getFileAccess().getOutputStream(outInfo.getFile());
				try {
					OutputStream channel = out.getChannel();
					stream.writeTo(channel);
					channel.flush();
				} catch (IOException e) {
					out.reportProblem();
					String msg = "Problems writing to channel " + outInfo;
					getLog().error(msg,e);
					throw new CurtainUncheckedException(msg,e);
				}finally{
					FileIO.closeQuietly(out);
				}				
			}
		}	
		mapper.reset();
	}

	private void flushAndClose(FileInfo file, FileChannelAccess<PrintWriter> outWriter) {
		if(null != outWriter){
			if(outWriter.isOpen()){
				try {
					outWriter.getChannel().flush();
				} catch (IOException e) {
					getLog().error("Problems flushing channel " + file);
				}
			}
			FileIO.closeQuietly(outWriter);			
		}
	}

//	protected EntryMapper getMapper(FileInfo outFile, FileChannelAccess<PrintWriter> out) {
//		return new MyEntryMapper(outFile.getFileType().getWriter(),out);
//	}

	private void _throwException(String msg, FileInfo inFile,FileInfo outFile, Long currPos, IOException e) throws CurtainUncheckedException {
		throw new CurtainUncheckedException(
				msg 
				+ "; Pos "
				+ currPos 
				+ "; from " 
				+ inFile
				+ "; to "
				+ outFile,
				e);
	}

	private FileInfo _getInputFileInfo(Integer inputId) {
		return inputId2FileIdx.get(inputId);
	}

	private FileInfo _getOutputFileInfo(Integer outputId) {
		return outputId2FileIdx.get(outputId);
	}	

	/**
	 * {@inheritDoc}
	 */
	public void addContigPosition(Integer fileId, Long pos) {
		_writeEntry(fileId,pos);
	}

	/**
	 * {@inheritDoc}
	 */
	public void addRead(Integer fileId, String key, Strand strand, Long position) {
		_writeEntry(fileId, position);
	}

	/**
	 * {@inheritDoc}
	 */
	public void addReadBatch(Integer[] fileIds, String[] keys,Strand[] strands, Long[] positions) {
		_writeBatch(fileIds,positions);
	}

	/**
	 * {@inheritDoc}
	 */
	public void addReadPosition(Integer fileId, Strand strand, Long pos) {
		_writeEntry(fileId,pos);
	}

	/**
	 * {@inheritDoc}
	 */
	public void openWrite(boolean append) throws IOException {
		if(!append){
			throw new CurtainUncheckedException("Only appending to output files supported at the moment!!!");
		}
	}

	/**
	 * {@inheritDoc}
	 */
	public Integer registerFile(FileInfo fileInfo, Integer category) {
		_addInputFile(fileInfo, category);
		Integer inId = _getInputfileId(fileInfo);
		return inId;
	}

	/**
	 * {@inheritDoc}
	 */
	public void synchronize(FileIndexReader reader) {
		throw new NotImplementedException("Not supported!");
	}

	/**
	 * {@inheritDoc}
	 */
	public void addContig(FileInfo fi, String key, Long pos) {
		_writeEntry(fi, pos);
	}

	/**
	 * {@inheritDoc}
	 */
	public void addContigPosition(FileInfo fi, Long pos) {
		_writeEntry(fi, pos);
	}

	/**
	 * {@inheritDoc}
	 */
	public void addFile(FileInfo fi, Integer category) {
		_addInputFile(fi, category);
	}
	
	private void _addInputFile(FileInfo fi, Integer category){
		Integer fId = _getFileId(fi);
		String strIdx = _buildOutputIdx(fi, category);
		Integer outId = category2OutfileIdx.get(strIdx);
		in2OutIdIdx.put(fId, outId);
	}

	private Integer _getFileId(FileInfo fi) {
		Integer fId = this.inputFile2IdIdx.get(fi);
		if(null == fId){
			/* For the moment throw exception*/
			throw new CurtainUncheckedException("Please add File id for File first before using!!!");
//			fId = _nextFildId();
//			_addInputFileId(fId, fi);
		}
		return fId;
	}

	private Integer iOutId = 0;
	
	private Integer _nextOutputFildId() {
		Integer fId = iOutId++;
		while(null != _getOutputFileInfo(fId)){
			fId = iOutId++;
		}
		return fId;
	}

	/**
	 * {@inheritDoc}
	 */
	public void addFileId(Integer fid, FileInfo fi) {
		_addInputFileId(fid, fi);
	}
	
	private void _addInputFileId(Integer fid,FileInfo fi){
		this.inputFile2IdIdx.put(fi, fid);
		this.inputId2FileIdx.put(fid, fi);
	}

	/**
	 * {@inheritDoc}
	 */
	public void addInsertLength(Integer category, Integer insLen) {
		/* ignore */
	}

	/**
	 * {@inheritDoc}
	 */
	public void addInsertLengthSd(Integer category, Integer sd) {
		/* ignore */
	}

	/**
	 * {@inheritDoc}
	 */
	public void addRead(FileInfo fi, String key, Strand strand, Long pos) {
		_writeEntry(fi, pos);
	}

	/**
	 * {@inheritDoc}
	 */
	public void addReadPosition(FileInfo fi, Strand strand, Long pos) {
		_writeEntry(fi, pos);
	}

	/**
	 * {@inheritDoc}
	 */
	public void close() throws IOException {
		/* Do nothing for the moment TODO */
	}
	
	private FileAccessService getFileAccess(){
		return fas;
	}
	
	public static class MyEntryMapper implements EntryMapper{

		private final Log log = LogFactory.getLog(this.getClass());
		private final AtomicReference<WriteUtil> writer = new AtomicReference<WriteUtil>();
		private final AtomicReference<FileChannelAccess<PrintWriter>> out = new AtomicReference<FileChannelAccess<PrintWriter>>();
		private final AtomicReference<IdentifierProcessor> from = new AtomicReference<IdentifierProcessor>();
		private final AtomicReference<IdentifierProcessor> to = new AtomicReference<IdentifierProcessor>();

		/**
		 * This mapper does NOT close / flush the {@link FileChannelAccess} {@link PrintWriter}
		 * @param buffered
		 */
		public MyEntryMapper() {
		}

		public void mapEntry(EntryMapperBean bean) {
			FileChannelAccess<PrintWriter> outLocal = getOutWriter();
			try{
				PrintWriter pw = outLocal.getChannel();
				if(getLog().isTraceEnabled()){
					getLog().trace("Write bean " + bean.getId() + " --> " + pw);
				}
				getWriter().write(pw, process(bean));
			} catch (Exception e) {
				outLocal.reportProblem();
				throw new CurtainUncheckedException(e);
			}
		}

		private WriteUtil getWriter() {
			return writer.get();
		}
		
		public void setWriter(WriteUtil writer) {
			this.writer.set(writer);
		}
		
		public void setChannelAccess(FileChannelAccess<PrintWriter> out) {
			this.out.set(out);
		}
		
		private Log getLog() {
			return log;
		}
		
		private EntryMapperBean process(EntryMapperBean bean){
			EntryMapperBean newBean = bean;
			if(null != getFrom() && null != getTo()){
				newBean = new EntryMapperBean(bean);
				Strand strand = getFrom().getStrand(newBean.getId());
				newBean.setId(
						getTo().addStrandInformation(
						getFrom().removeStrandInformation(newBean.getId()), strand));
			}
			return newBean;
		}

		private IdentifierProcessor getFrom() {
			return from.get();
		}
		
		private IdentifierProcessor getTo() {
			return to.get();
		}
		
		public void setIdProcessorFrom(IdentifierProcessor from) {
			this.from.set(from);
		}
		public void setIdProcessorTo(IdentifierProcessor to) {
			this.to.set(to);
		}
		public FileChannelAccess<PrintWriter> getOutWriter() {
			return out.get();
		}			
		
	}
	
	public static class BufferedEntryMapper implements EntryMapper{

		private final Log log = LogFactory.getLog(this.getClass());
		private final AtomicLong byteCount = new AtomicLong(0);
		private final AtomicReference<WriteUtil> writer = new AtomicReference<WriteUtil>();
		private final ConcurrentMap<Integer, ByteArrayOutputStream> streamMap = UtilsCollect.newConcurrentMap();
		private final ConcurrentMap<Integer, PrintWriter> printMap = UtilsCollect.newConcurrentMap();
		private final AtomicInteger outId = new AtomicInteger(-1);
		private final AtomicReference<IdentifierProcessor> from = new AtomicReference<IdentifierProcessor>();
		private final AtomicReference<IdentifierProcessor> to = new AtomicReference<IdentifierProcessor>();

		/**
		 * This mapper does NOT close / flush the {@link FileChannelAccess} {@link PrintWriter}
		 * @param buffered
		 */
		public BufferedEntryMapper() {
		}

		public void setOutputId(Integer outId) {
			this.outId.set(outId);
		}

		public long getByteCount() {
			return byteCount.get();
		}

		public void mapEntry(EntryMapperBean bean) {
			Integer id = getOutId();
			EntryMapperBean outBean = process(bean);
			if(id < 0){
				throw new CurtainUncheckedException("Output id is not set!!! " + id + " for file " + bean.getInfo() + "; bean " + bean);
			}
			addBean(id,outBean);			
		}

		private void addBean(Integer id, EntryMapperBean outBean) {
			PrintWriter pw = getPrintWriter(id);
			ByteArrayOutputStream stream = getStream(id);
			int start = stream.size();
			getWriter().write(pw, outBean);
			pw.flush();
			int end = stream.size();
			byteCount.addAndGet(end-start);
		}

		private PrintWriter getPrintWriter(Integer id) {
			PrintWriter pw = this.printMap.get(id);
			if(null == pw){
				this.printMap.putIfAbsent(id, new PrintWriter(getStream(id)));
				pw = this.printMap.get(id);
			}
			return pw;
		}
		
		private ByteArrayOutputStream getStream(Integer id) {
			ByteArrayOutputStream stream = this.streamMap.get(id);
			if(null == stream){
				streamMap.putIfAbsent(id, new ByteArrayOutputStream());
				stream = streamMap.get(id);
			}
			return stream;
		}
		
		public ConcurrentMap<Integer, ByteArrayOutputStream> getStreamMap() {
			return streamMap;
		}
		
		public ConcurrentMap<Integer, PrintWriter> getPrintMap() {
			return printMap;
		}
		
		public void reset() {
			this.byteCount.set(0);
			this.printMap.clear();
			this.streamMap.clear();
		}
		
		private Integer getOutId() {
			return outId.get();
		}

		private WriteUtil getWriter() {
			return writer.get();
		}
		
		public void setWriter(WriteUtil writer) {
			this.writer.set(writer);
		}
		
		private Log getLog() {
			return log;
		}
		
		private EntryMapperBean process(EntryMapperBean bean){
			EntryMapperBean newBean = bean;
			if(null != getFrom() && null != getTo()){
				newBean = new EntryMapperBean(bean);
				Strand strand = getFrom().getStrand(newBean.getId());
				newBean.setId(
						getTo().addStrandInformation(
						getFrom().removeStrandInformation(newBean.getId()), strand));
			}
			return newBean;
		}

		private IdentifierProcessor getFrom() {
			return from.get();
		}
		
		private IdentifierProcessor getTo() {
			return to.get();
		}
		
		public void setIdProcessorFrom(IdentifierProcessor from) {
			this.from.set(from);
		}
		public void setIdProcessorTo(IdentifierProcessor to) {
			this.to.set(to);
		}	
		
	}

	@Override
	public void addCategory(Integer category, String id) {
		throw new NotImplementedException();
	}
	
}
